import os, sys

parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 5)))
sys.path.insert(0, parent_path)

# from queue import Queue
from core.task.modules.processors.base_processor import BaseProcessor
from threading import Thread
from typing import *
import cv2, time, traceback
import numpy as np
from core.utils.logger import logging
from queue import Queue


class SteamReader(BaseProcessor):
    """任务管道的视频预处理组件
    跟随任务启动后, 会在task_init时指定视频地址"""

    def __init__(
        self,
        keys,
        SteamReader_path=None,
        SteamReader_size=None,
    ) -> None:
        """
        input:
            keys:{
                out: 在流转数据box中的键名
                finish_tag: 文件读取完成后向databox中写入的key
            }
            VideoPreprocessor_path: 视频地址
            VideoPreprocessor_path_size: 视频图像的大小
        """
        super().__init__(keys)
        self.keys = keys
        self.steamreader_path = SteamReader_path
        self.steamreader_size = SteamReader_size
        self.cap = None
        self.frame = None
        self.frame_num = 0
        self.run = None

    def _init_check(self):
        assert "out" in self.keys
        assert "finish_tag" in self.keys
        assert not self.steamreader_path is None
        assert not self.cap is None
        assert not self.steamreader_size is None

    def _set(self, data: Dict):
        """
        当本类作为module被Task Pipeliner调用时起效, 作为任务初始化函数被调用
        """
        if data.get("SteamReader_path", False):
            self.steamreader_path = data["SteamReader_path"]
        if data.get("SteamReader_size", False):
            self.steamreader_size = data["SteamReader_size"]
        self.cap = self.init_cap(self.steamreader_path)

    def init_cap(self, steamreader_path):
        cap = cv2.VideoCapture(steamreader_path)
        assert cap.isOpened(), "{} can't open".format(steamreader_path)
        ret, frame = cap.read()
        assert ret, "{} can't open".format(steamreader_path)
        self.frame = frame
        self.frame_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        return cap

    def _call(self, data: Dict) -> Dict:
        data[self.keys["out"]] = self.frame[None, ...]
        # 二次判断
        ret, frame = self.cap.read()
        if not ret:
            ret, frame = self.cap.read()

        # 二次判断失败, 中断逻辑处理
        if not ret:
            if (
                int(self.cap.get(cv2.CAP_PROP_POS_FRAMES)) >= self.frame_num
            ):  # 视频已经读取完整
                logging.info("SteamReader read done.")
            else:
                logging.warning("SteamReader.cap.read() fail in __call__(), retry ....")
                for _ in range(4):  # 重试4次
                    ret, frame = self.cap.read()  # 更新读取结果
                    if ret:
                        logging.warning("SteamReader.cap.read() retry success!")
                        break
                if not ret:
                    logging.error("SteamReader.cap.read() retry FALSE, close Task")

        # 最终逻辑判断, 如果仍为False则退出
        if not ret:
            data[self.keys["finish_tag"]] = True
            logging.info("SteamReader close.")
            self.close()
        else:
            self.frame = cv2.resize(frame, self.steamreader_size)

        return data


if __name__ == "__main__":
    steam_reader = SteamReader(
        keys={"out": "frame_data", "finish_tag": "read_done"},
        SteamReader_path=None,
        SteamReader_size=[1920, 1080],
    )
    save_path = "/home/smartgis/workspace/project/smart_server/data/output/temp.mp4"
    data = {
        "SteamReader_path": "/home/smartgis/workspace/data/原视频.mp4",
        "SteamReader_size": [1920, 1080],
    }
    steam_reader.set(data)
    steam_reader.start()
    steam_reader.cap.set(cv2.CAP_PROP_POS_FRAMES, 10800)
    idx = 0
    video_write = None
    while 1:
        idx += 1
        data = steam_reader(data)
        if video_write is None:
            video_write = cv2.VideoWriter(
                save_path, cv2.VideoWriter_fourcc(*"XVID"), 30, (1920, 1080)
            )
        video_write.write(data["frame_data"][0])
        if idx > (2 * 60 * 30):
            break
        if data.get("read_done", False):
            break
        # time.sleep(1)
    video_write.release()
    steam_reader.close()
